package cn.jtcoding.web.framework.http;

import lombok.Getter;
import lombok.Setter;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;

/**
 * 该类继承HttpServletResponseWrapper，先获取到回复的内容，再将内容重新set到outputStream实现回复体多读。
 */
public class MultiReadHttpServletResponse extends HttpServletResponseWrapper {

	private HttpServletResponse originalResponse = null;
	private PrintWriter writer = null;
	private final WapperedServletOutputStream wrappedServletOutputStream;
	@Setter
	@Getter
	private String remoteAddr;
	@Setter
	@Getter
	private String localAddr;

	public MultiReadHttpServletResponse(HttpServletResponse response) throws IOException {
		super(response);
		this.originalResponse = response;
		OutputStream outputStream = response.getOutputStream();
		this.wrappedServletOutputStream = new WapperedServletOutputStream(outputStream);
		this.writer = new PrintWriter(new OutputStreamWriter(wrappedServletOutputStream, this.getCharacterEncoding()),
				true);
	}

	public void reWriteOutputStream(byte[] bytes) throws IOException {
		wrappedServletOutputStream.write(bytes);
	}

	@Override
	public ServletOutputStream getOutputStream() throws IOException {
		return wrappedServletOutputStream;
	}

	@Override
	public PrintWriter getWriter() throws IOException {
		return writer;
	}

	@Override
	public void flushBuffer() throws IOException {
		if (wrappedServletOutputStream != null) {
			wrappedServletOutputStream.flush();
		}
		if (writer != null) {
			writer.flush();
		}
	}

	public String getResponseBody() throws IOException {
		this.flushBuffer();// 将out、writer中的数据强制输出到WapperedResponse的buffer里面
		byte[] bytes = wrappedServletOutputStream.toByteArray();
		return new String(bytes, this.getCharacterEncoding());
	}

	/**
	 * 获取最原始的HttpServletResponse
	 *
	 * @return
	 */
	public HttpServletResponse getOriginalResponse() {
		return originalResponse;
	}

	private static class WapperedServletOutputStream extends ServletOutputStream {

		private OutputStream stream = null;
		private final ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(1024);

		public WapperedServletOutputStream(OutputStream stream) throws IOException {
			this.stream = stream;
		}

		@Override
		public void write(int b) throws IOException {
			stream.write(b);
			byteArrayOutputStream.write(b);
		}

		@Override
		public boolean isReady() {
			// TODO Auto-generated method stub
			return false;
		}

		@Override
		public void setWriteListener(WriteListener listener) {
			// TODO Auto-generated method stub

		}

		public byte[] toByteArray() {
			return byteArrayOutputStream.toByteArray();
		}
	}

}
